#pragma once

namespace Ut
{
	template <typename T>
	class WeakPtr;

	template<typename T>
	void Destroy(void* p)
	{
		((T*) p)->~T();
		delete p;
	}

	typedef void (*DestroyFunctionPtr)(void* p);

	struct ReferenceCountedPointer
	{
		template <typename T>
		ReferenceCountedPointer(T* p, DestroyFunctionPtr f = &Destroy<T>) : m_p(p), m_f(f), count(0), weakCount(0) {}

		~ReferenceCountedPointer()
		{
			Assert(m_p == 0);
		}

		template <typename T>
		void Release()
		{
			Assert(count == 0);

			if (m_p)
			{
				void* p = m_p;
				m_p = 0; // In case the object has a pointer to itself.
				m_f(p);
			}
		}

		template <typename T>
		T* Pointer() {return (T*) m_p;}

		void* m_p;
		DestroyFunctionPtr m_f;
		int count;
		int weakCount;
	};

	template <typename T>
	class SharedPtr
	{
	public:
		typedef T* PointerType;

		SharedPtr() : m_p(0) {}

		explicit SharedPtr(typename PointerType p)
		{
			if (p)
			{
				m_p = new ReferenceCountedPointer(p);
				m_p->count = 1;
				m_p->weakCount = 0;
			}
			else
			{
				m_p = 0;
			}
		}

		SharedPtr(const SharedPtr& other)
		{
			m_p = other.m_p;
			if (m_p) m_p->count++;
		}

		SharedPtr(SharedPtr&& other)
		{
			m_p = other.m_p;
			other.m_p = 0;
		}

		explicit SharedPtr(const WeakPtr<T>& other);

		SharedPtr& operator = (const SharedPtr& other)
		{
			Release();
			m_p = other.m_p;
			if (m_p) m_p->count++;
			return *this;
		}

		SharedPtr& operator = (SharedPtr&& other)
		{
			Release();
			m_p = other.m_p;
			other.m_p = 0;
			return *this;
		}

		void SetReference(ReferenceCountedPointer* ref)
		{
			Release();
			m_p = ref;
			if (m_p) m_p->count++;
		}

		~SharedPtr()
		{
			Release();
		}

		typename PointerType operator -> () 
		{
			Assert(m_p);
			return (PointerType) *this;
		}

		typename const PointerType operator -> () const
		{
			Assert(m_p);
			return (PointerType) *this;
		}

		operator typename PointerType () const
		{
			return m_p ? m_p->Pointer<T>() : 0;
		}

		void Release()
		{
			if (m_p)
			{
				m_p->count--;

				if (m_p->count == 0)
				{
					int weak = m_p->weakCount;

					m_p->Release<T>(); // If weak > 0, this might delete m_p.

					if (weak == 0)
						delete m_p;
				}

				m_p = 0;
			}
		}

		void Swap(SharedPtr& other)
		{
			std::swap(m_p, other.m_p);
		}
		
	private:
		template <typename U>
		friend class WeakPtr;

		ReferenceCountedPointer* m_p;
	};

	template <typename T>
	class WeakPtr
	{
	public:
		typedef T* PointerType;

		WeakPtr() : m_p(0) {}

		WeakPtr(const SharedPtr<T>& other)
		{
			m_p = other.m_p;
			if (m_p) m_p->weakCount++;
		}

		WeakPtr(const WeakPtr& other)
		{
			m_p = other.m_p;
			if (m_p) m_p->weakCount++;
		}

		WeakPtr(WeakPtr&& other)
		{
			m_p = other.m_p;
			other.m_p = 0;
		}

		WeakPtr& operator = (const SharedPtr<T>& other)
		{
			Release();
			m_p = other.m_p;
			if (m_p) m_p->weakCount++;
			return *this;
		}

		template <typename U>
		WeakPtr& operator = (const SharedPtr<U>& other)
		{
			Release();
			m_p = other.m_p;
			if (m_p) m_p->weakCount++;
			return *this;
		}

		WeakPtr& operator = (const WeakPtr& other)
		{
			Release();
			m_p = other.m_p;
			if (m_p) m_p->weakCount++;
			return *this;
		}

		WeakPtr& operator = (WeakPtr&& other)
		{
			Release();
			m_p = other.m_p;
			other.m_p = 0;
			return *this;
		}

		void SetReference(ReferenceCountedPointer* ref)
		{
			Release();
			m_p = ref;
			if (m_p) m_p->weakCount++;
		}

		~WeakPtr()
		{
			Release();
		}

		explicit operator bool () const
		{
			return m_p ? (m_p->Pointer<T>() != 0) : false;
		}

		SharedPtr<T> Lock() const
		{
			return SharedPtr<T>(*this);
		}

		bool operator == (const WeakPtr& other) const
		{
			return (m_p ? m_p->Pointer<T>() : 0) == (other.m_p ? other.m_p->Pointer<T>() : 0);
		}

		bool operator != (const WeakPtr& other) const
		{
			return (m_p ? m_p->Pointer<T>() : 0) != (other.m_p ? other.m_p->Pointer<T>() : 0);
		}

		void Release()
		{
			if (m_p)
			{
				m_p->weakCount--;

				if (m_p->count == 0)
				{
					if (m_p->weakCount == 0)
						delete m_p;
				}
			}

			m_p = 0;
		}
		
	private:
		friend class SharedPtr<T>;

		ReferenceCountedPointer* m_p;
	};

	template <typename T>
	inline SharedPtr<T>::SharedPtr(const WeakPtr<T>& other)
	{
		m_p = other.m_p;
		if (m_p) m_p->count++;
	}

	template <typename T>
	inline SharedPtr<T> MakeShared(T* p)
	{
		return SharedPtr<T>(p);
	}

	template <typename T>
	inline void Swap(SharedPtr<T>& a, SharedPtr<T>& b)
	{
		a.Swap(b);
	}
}